/*
* Copyright (C) Microsoft. All rights reserved.  
*
* Licensed under the Apache License, Version 2.0 (the "License"); you may not 
* use this file except in compliance with the License.  You may obtain a copy 
* of the License at http://www.apache.org/licenses/LICENSE-2.0  
* 
* THIS CODE IS PROVIDED *AS IS* BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, EITHER EXPRESS OR IMPLIED, INCLUDING WITHOUT LIMITATION ANY IMPLIED 
* WARRANTIES OR CONDITIONS OF TITLE, FITNESS FOR A PARTICULAR PURPOSE, 
* MERCHANTABLITY OR NON-INFRINGEMENT. 
*
* See the Apache License, Version 2.0 for specific language governing 
* permissions and limitations under the License.
*
*/
#pragma once

// NOTE:
// ideally I'd make xml_helper an interface and then have a win32 implementation for it
// think about doing that if it is worth the extra work at some point

#include <Windows.h>
#include <wrl.h>
#include <xmllite.h>

#pragma comment(lib, "xmllite.lib")

#include <stack>

#include <baja/common.hpp>
#include <baja/error.hpp>
#include <baja/assert.hpp>

#include <baja/stream/stream.hpp>
#include <baja/stream/win32/details/win32_stream_wrapper.hpp>

namespace baja { namespace win32 { namespace details {

namespace WRL = Microsoft::WRL;
namespace w32d = ::baja::win32::details;

struct xml_element
{
    std::wstring name;
    std::wstring namespaceUri;
    std::wstring text;
    std::vector<xml_element> children;
    std::map<std::wstring, std::wstring> attributes;
};

// xml_helper class
class xml_helper
{
public:
    xml_helper() {}

    void load(
        const wchar* xml,
        const uint32 size,
        xml_element* root
        ) const; 

    void load(
        std::shared_ptr<istream>& fileStream, 
        xml_element* root
        ) const;

    xml_element* findElementByName(
        const wchar* elementName,
        xml_element* rootOfSearch
        ) const;

    xml_element* findElementByAttributeValue(
        const wchar* attributeName,
        const wchar* attributeValue,
        xml_element* rootOfSearch
        ) const;

    void save(
        std::shared_ptr<istream>& saveStream,
        xml_element* root,
        const wchar* encoding = nullptr
        ) const;

    void save(
        const std::wstring& saveString,
        const xml_element* root
        ) const;
};

// implementation

// loads XML from the specified string
inline void xml_helper::load(
    const wchar* xml,
    const size_t size,
    xml_element* root
    ) const
{
    BAJA_CHECK_ARG(xml != nullptr);
    BAJA_CHECK_ARG(size > 0);
    BAJA_CHECK_ARG(root != nullptr);

    std::shared_ptr<imemory_stream> stream;
    stream = createMemoryStream(reinterpret_cast<const byte*>(xml), size * sizeof(wchar));

    this->load(std::dynamic_pointer_cast<istream>(stream), root);
}

// loads XML from the specified stream
inline void xml_helper::load(
    std::shared_ptr<istream>& fileStream, 
    xml_element* root
    ) const
{
    BAJA_CHECK_ARG(fileStream);
    BAJA_CHECK_ARG(root);

    win32_result wr;

    // create the reader
    WRL::ComPtr<IXmlReader> reader;
    wr = CreateXmlReader(__uuidof(IXmlReader), (void**)&reader, NULL);
    WRL::ComPtr<IStream> streamWrapper = reinterpret_cast<IStream*>(win32_stream_wrapper::create(fileStream));
    wr = reader->SetInput(streamWrapper.Get());

    // initialize the element stack
    std::stack<xml_element*> elementStack;
    elementStack.push(root);
    uint32 currentDepth = 0;

    // NOTE: XmlLite likes to return S_FALSE
    HRESULT hr = S_OK;
    while (hr == S_OK)
    {
        // read the node type from XML
        XmlNodeType nodeType = XmlNodeType_None;
        if (hr == S_OK)
        {
            hr = reader->Read(&nodeType);
        }

        if (hr == S_OK)
        {
            // get current depth of XML
            uint32 newDepth = 0;
            if (hr == S_OK)
            {
                hr = reader->GetDepth(&newDepth);
            }

            if (hr == S_OK)
            {
                // is node an element type?
                if (nodeType == XmlNodeType_Element)
                {
                    // build up a new element
                    xml_element newElement;

                    const wchar* name = NULL;
                    if (hr == S_OK)
                    {
                        hr = reader->GetLocalName(&name, NULL);
                    }
                    if (hr == S_OK)
                    {
                        // set name of new element
                        newElement.name = name;

                        // load all of the element's attributes
                        if (hr == S_OK)
                        {
                            hr = reader->MoveToFirstAttribute();
                        }
                        while (hr == S_OK) 
                        {
                            // load attribute name
                            const wchar* attr = NULL;
                            if (SUCCEEDED(hr))
                            {
                                hr = reader->GetLocalName(&attr, NULL);
                            }
                        
                            // load attribute value
                            const wchar* val = NULL;
                            if (SUCCEEDED(hr))
                            {
                                hr = reader->GetValue(&val, NULL);
                            }

                            // store the attribute
                            if (SUCCEEDED(hr))
                            {
                                newElement.attributes[attr] = val;
                            }

                            // move to next attribute
                            if (SUCCEEDED(hr))
                            {
                                hr = reader->MoveToNextAttribute();
                            }
                        }

                        // reset HR since S_FALSE is returned
                        if (SUCCEEDED(hr))
                        {
                            hr = S_OK;
                        }

                        // pop element stack as needed
                        if (SUCCEEDED(hr))
                        {
                            // adjust stack to correct depth
                            BAJA_VERIFY(currentDepth >= newDepth);
                            while (currentDepth > newDepth)
                            {
                                elementStack.pop();
                                currentDepth--;
                            }

                            // add element as a child of element at top of stack
                            xml_element* top = elementStack.top();
                            top->children.push_back(newElement);

                            // push new element onto stack
                            xml_element* elementToPush = &(top->children[top->children.size()-1]);
                            elementStack.push(elementToPush);
                            currentDepth++;
                        }
                    }
                }
                else if (nodeType == XmlNodeType_Text || nodeType == XmlNodeType_CDATA)
                {
                    const wchar* text = NULL;
                    if (hr == S_OK)
                    {
                        hr = reader->GetValue(&text, NULL);
                        xml_element* ele = elementStack.top();
                        ele->text = text;
                    }
                }
                else if (nodeType == XmlNodeType_EndElement)
                {
                    // an element has ended, adjust stack to correct depth
                    newDepth--;
                    BAJA_VERIFY(currentDepth >= newDepth);
                    while (currentDepth > newDepth)
                    {
                        elementStack.pop();
                        currentDepth--;
                    }
                }
            }
        }
    }

    if (SUCCEEDED(hr))
    {
        // reset to S_OK since S_FALSE is being used by XmlLite
        hr = S_OK;
    }
    wr = hr; 
}

// finds an element by its name
inline xml_element* xml_helper::findElementByName( 
    const wchar* elementName, 
    xml_element* rootOfSearch
    ) const
{
    xml_element* foundElement = NULL;

    if (elementName != NULL && rootOfSearch != NULL)
    {
        if (rootOfSearch->name == elementName)
        {
            foundElement = rootOfSearch;
        }
        else
        {
            for (uint32 i = 0; i < rootOfSearch->children.size() && foundElement == NULL; i++)
            {
                xml_element* childToSearch = &(rootOfSearch->children[i]);
                foundElement = this->findElementByName(elementName, childToSearch);
            }
        }
    }

    return foundElement;
}

//
// finds first element that has a matching attribute and attribute value
//
inline xml_element* xml_helper::findElementByAttributeValue( 
    const wchar* attributeName, 
    const wchar* attributeValue, 
    xml_element* rootOfSearch
    ) const
{
    xml_element* foundElement = NULL;

    if (attributeName != NULL && attributeValue != NULL && rootOfSearch != NULL)
    {
        if (rootOfSearch->attributes[attributeName] == attributeValue)
        {
            foundElement = rootOfSearch;
        }
        else
        {
            for (uint32 i = 0; i < rootOfSearch->children.size() && foundElement == NULL; i++)
            {
                xml_element* childToSearch = &(rootOfSearch->children[i]);
                foundElement = this->findElementByAttributeValue(attributeName, attributeValue, childToSearch);
            }
        }
    }

    return foundElement;
}

/*
NYI:
//
// saves XML into the specified stream
//
void 
xml_helper::save(
    istream* saveStream,
    xml_element* root,
    wchar* encoding
    ) const
{
    // include the byte order mark because we are writing to stream
    bool includeByteOrderMark = true;
    return SaveAsXml(saveStream, root, encoding, includeByteOrderMark);
}

//
// saves XML into the specified string
//
void 
xml_helper::save(
    stringw* saveString,
    Xml_element* root
    ) const
{
    E_POINTER_Return(saveString);

    Auto<IMemoryStream> memoryStream;
    TRACEHR hr = CreateAllocatedMemoryStream(NULL, 0, StreamFlagsGrowable, &memoryStream);

    if (SUCCEEDED(hr))
    {
        wchar* encoding = L"utf-16";
        bool includeByteOrderMark = false; // writing to string, so do not include the byte order mark
        hr = SaveAsXml(memoryStream, root, encoding, includeByteOrderMark);
    }

    if (SUCCEEDED(hr))
    {
        // leave space for the trailing NULL character
        uint32 dataSizeInBytes = memoryStream->GetBufferSize() + sizeof(wchar); 
        uint8* data = new uint8[dataSizeInBytes];
        hr = HrFromAlloc(data);

        if (SUCCEEDED(hr))
        {
            ZeroMemory(data, dataSizeInBytes);
            CopyMemory(data, memoryStream->GetBuffer(), memoryStream->GetBufferSize());
            (*saveString) = (wchar*)data;
        }
    }

    return hr;
}
*/

}}} // namespace baja.win32.details
